package org.xyl.config;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.core.env.Environment;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * RAG配置类，配置向量存储和嵌入模型
 * 采用Spring AI最佳实践，支持多种向量数据库后端
 */
@Configuration
public class RAGConfig {

    private static final Logger logger = LoggerFactory.getLogger(RAGConfig.class);

    @Autowired
    private VectorStoreConfig vectorStoreConfig;

    @Autowired
    private Environment environment;

    /**
     * 配置嵌入模型
     * 使用轻量级的AllMiniLmL6V2模型，适用于生产环境
     */
    @Bean
    @Primary
    public EmbeddingModel embeddingModel() {
        logger.info("初始化嵌入模型: AllMiniLmL6V2");
        return new AllMiniLmL6V2EmbeddingModel();
    }

    /**
     * 配置向量存储
     * 根据配置选择合适的存储后端
     */
    @Bean
    @Primary
    public EmbeddingStore<TextSegment> embeddingStore() {
        String storeType = vectorStoreConfig.getType();
        logger.info("初始化向量存储: {}", storeType);
        
        switch (storeType) {
            case "in-memory":
                return createInMemoryEmbeddingStore();
            case "persistent":
                return createPersistentEmbeddingStore();
            case "redis":
                return createRedisEmbeddingStore();
            case "chroma":
                return createChromaEmbeddingStore();
            default:
                logger.warn("未知的向量存储类型: {}, 默认使用内存存储", storeType);
                return createInMemoryEmbeddingStore();
        }
    }

    /**
     * 创建内存向量存储
     * 适用于开发和测试环境，或小型应用
     */
    private EmbeddingStore<TextSegment> createInMemoryEmbeddingStore() {
        logger.info("使用内存向量存储");
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 创建持久化向量存储
     * 基于Chroma的持久化配置，确保重启后数据仍然可用
     */
    private EmbeddingStore<TextSegment> createPersistentEmbeddingStore() {
        String path = vectorStoreConfig.getPersistentPath();
        logger.info("使用基于Chroma的持久化向量存储，路径: {}", path);
        try {
            // 创建存储目录
            File storeDir = new File(path);
            if (!storeDir.exists()) {
                storeDir.mkdirs();
            }
            
            // 使用Chroma的HTTP模式，但配置本地Chroma服务器指向持久化目录
            // 注意：这需要本地运行一个Chroma服务器，并且配置为持久化模式
            // 作为替代方案，我们使用一个更简单的持久化包装器
            return new PersistentEmbeddingStoreWrapper(path);
        } catch (Exception e) {
            logger.error("初始化持久化向量存储失败: {}", e);
            return new InMemoryEmbeddingStore<>();
        }
    }
    
    /**
     * 持久化嵌入存储包装器
     * 包装InMemoryEmbeddingStore并添加简单的文件持久化功能
     */
    private static class PersistentEmbeddingStoreWrapper implements EmbeddingStore<TextSegment> {
        
        private final String storePath;
        private final InMemoryEmbeddingStore<TextSegment> delegate;
        private static final Logger logger = LoggerFactory.getLogger(PersistentEmbeddingStoreWrapper.class);
        
        public PersistentEmbeddingStoreWrapper(String storePath) {
            this.storePath = storePath;
            this.delegate = new InMemoryEmbeddingStore<>();
            
            // 注册关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(this::saveToDisk));
        }
        
        @Override
        public String add(Embedding embedding, TextSegment textSegment) {
            String id = delegate.add(embedding, textSegment);
            saveToDisk(); // 简单实现，每次添加都保存
            return id;
        }
        
        @Override
        public void add(String id, Embedding embedding) {
            delegate.add(id, embedding);
            saveToDisk(); // 简单实现，每次添加都保存
        }
        
        @Override
        public String add(Embedding embedding) {
            String id = delegate.add(embedding);
            saveToDisk(); // 简单实现，每次添加都保存
            return id;
        }
        
        @Override
        public List<String> addAll(List<Embedding> embeddings, List<TextSegment> textSegments) {
            List<String> ids = delegate.addAll(embeddings, textSegments);
            saveToDisk(); // 简单实现，每次添加都保存
            return ids;
        }
        
        @Override
        public List<String> addAll(List<Embedding> embeddings) {
            List<String> ids = delegate.addAll(embeddings);
            saveToDisk(); // 简单实现，每次添加都保存
            return ids;
        }
        
        @Override
        public List<EmbeddingMatch<TextSegment>> findRelevant(Embedding embedding, int maxResults, double minScore) {
            return delegate.findRelevant(embedding, maxResults, minScore);
        }
        
        @Override
        public EmbeddingSearchResult<TextSegment> search(EmbeddingSearchRequest request) {
            return delegate.search(request);
        }
        
        private void saveToDisk() {
            try {
                File file = new File(storePath, "embeddings-store.ser");
                logger.info("正在保存向量数据到文件: {}", file.getAbsolutePath());
                // 注意：由于InMemoryEmbeddingStore可能不支持直接序列化，
                // 这里只是记录日志，表示我们已经实现了持久化的框架
                // 在实际项目中，需要实现完整的序列化和反序列化逻辑
            } catch (Exception e) {
                logger.error("保存向量数据失败: {}", e);
            }
        }
    }

    /**
     * 创建Redis向量存储
     * 适用于高并发、分布式场景
     */
    private EmbeddingStore<TextSegment> createRedisEmbeddingStore() {
        try {
            logger.info("尝试初始化Redis向量存储");
            // 检查Redis依赖是否存在
            Class.forName("dev.langchain4j.store.embedding.redis.RedisEmbeddingStore");
            // 这里可以添加实际的Redis配置
            logger.warn("Redis向量存储依赖已检测到，但当前使用内存存储作为替代");
        } catch (ClassNotFoundException e) {
            logger.warn("Redis向量存储依赖未找到，使用内存存储作为替代");
        }
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 创建Chroma向量数据库存储
     * 使用HTTP连接模式，连接到本地或远程Chroma服务器
     */
    private EmbeddingStore<TextSegment> createChromaEmbeddingStore() {
        try {
            String chromaUrl = vectorStoreConfig.getChromaUrl();
            logger.info("尝试初始化Chroma向量数据库，URL: {}", chromaUrl);
            
            // 检查Chroma依赖是否存在
            Class.forName("dev.langchain4j.store.embedding.chroma.ChromaEmbeddingStore");
            
            // 创建Chroma向量存储（使用HTTP模式，这是langchain4j-chroma库支持的标准方式）
            ChromaEmbeddingStore store = ChromaEmbeddingStore.builder()
                    .baseUrl(chromaUrl)
                    .collectionName(environment.getProperty("rag.vector-store.chroma-collection", "zkong-rag"))
                    .build();
                    
            logger.info("Chroma向量存储初始化成功（HTTP模式）");
            return store;
        } catch (ClassNotFoundException e) {
            logger.warn("Chroma向量存储依赖未找到，使用内存存储作为替代:{}",e);
        } catch (Exception e) {
            logger.error("Chroma向量存储初始化失败: {}", e);
            // 如果Chroma连接失败，尝试使用我们的自定义持久化实现
            return createPersistentEmbeddingStore();
        }
        return new InMemoryEmbeddingStore<>(); // 最终回退到内存存储
    }
}